#include "DHVideo.h"
#include "DHMfxEncode.h"

#include "d3d_device.h"
#include "d3d_allocator.h"
#include "d3d11_device.h"
#include "d3d11_allocator.h"
#include "sysmem_allocator.h"
#include "libyuv.h"

extern void test_start();
extern double test_check();

mfxStatus InitSession(MFXVideoSession& mfxSession)
{
    mfxInitParam initPar;
    mfxVersion version;
    mfxStatus sts = MFX_ERR_NONE;

    MSDK_ZERO_MEMORY(initPar);
    initPar.Version.Major = 1;
    initPar.Version.Minor = 0;
    initPar.GPUCopy =  0;
    initPar.Implementation = MFX_IMPL_HARDWARE_ANY;
    sts = mfxSession.InitEx(initPar);
    if(sts<MFX_ERR_NONE){
        initPar.Implementation = MFX_IMPL_SOFTWARE;
        sts = mfxSession.InitEx(initPar);
    }
    MSDK_CHECK_STATUS(sts, "m_mfxSession.InitEx failed");
    sts = MFXQueryVersion(mfxSession, &version);
    if (!CheckVersion(&version, MSDK_FEATURE_JPEG_ENCODE)) {
        msdk_printf(MSDK_STRING("error: Jpeg is not supported in the %d.%d API version\n"),version.Major, version.Minor);
        return MFX_ERR_UNSUPPORTED;
    }
    return sts;
}

mfxStatus InitMfxEncParams(mfxDHEncode* pEncode, mfxU32 CodecId,mfxU16 Width,mfxU16 Height, mfxU64 frameRate, mfxU16 bitRate)
{
    pEncode->mfxEncParams.mfx.CodecId = CodecId;
    pEncode->mfxEncParams.mfx.TargetUsage = MFX_TARGETUSAGE_2; // trade-off between quality and speed
    pEncode->mfxEncParams.mfx.RateControlMethod = MFX_RATECONTROL_CBR;
    pEncode->mfxEncParams.mfx.GopRefDist = 0;
    pEncode->mfxEncParams.mfx.GopPicSize = 0;
    pEncode->mfxEncParams.mfx.NumRefFrame = 0;
    pEncode->mfxEncParams.mfx.IdrInterval = 0;

    pEncode->mfxEncParams.mfx.CodecProfile = 0;
    pEncode->mfxEncParams.mfx.CodecLevel = 0;
    pEncode->mfxEncParams.mfx.InitialDelayInKB = 0;
    pEncode->mfxEncParams.mfx.GopOptFlag = 0;
    pEncode->mfxEncParams.mfx.BufferSizeInKB = 0;
    pEncode->mfxEncParams.mfx.TargetKbps = bitRate; // in Kbps
    pEncode->mfxEncParams.mfx.QPP = bitRate;
    pEncode->mfxEncParams.mfx.ICQQuality = bitRate;
    pEncode->mfxEncParams.mfx.LowPower = MFX_CODINGOPTION_OFF;
    pEncode->mfxEncParams.mfx.NumSlice = 0;
    pEncode->mfxEncParams.mfx.FrameInfo.FrameRateExtN = frameRate;
    pEncode->mfxEncParams.mfx.FrameInfo.FrameRateExtD = 1;
    pEncode->mfxEncParams.mfx.EncodedOrder = 0;
    // specify memory type
    if (D3D9_MEMORY == pEncode->memType)
        pEncode->mfxEncParams.IOPattern = MFX_IOPATTERN_IN_VIDEO_MEMORY;
    else
        pEncode->mfxEncParams.IOPattern = MFX_IOPATTERN_IN_SYSTEM_MEMORY;

    // frame info parameters
    pEncode->mfxEncParams.mfx.FrameInfo.FourCC = MFX_FOURCC_NV12;
    pEncode->mfxEncParams.mfx.FrameInfo.ChromaFormat = FourCCToChroma(MFX_FOURCC_NV12);
    pEncode->mfxEncParams.mfx.FrameInfo.PicStruct = MFX_PICSTRUCT_PROGRESSIVE;
    pEncode->mfxEncParams.mfx.FrameInfo.Shift = 0;

    // width must be a multiple of 16
    // height must be a multiple of 16 in case of frame picture and a multiple of 32 in case of field picture
    pEncode->mfxEncParams.mfx.FrameInfo.Width  = MSDK_ALIGN16(Width);
    pEncode->mfxEncParams.mfx.FrameInfo.Height = MSDK_ALIGN16(Height);

    pEncode->mfxEncParams.mfx.FrameInfo.CropX = 0;
    pEncode->mfxEncParams.mfx.FrameInfo.CropY = 0;
    pEncode->mfxEncParams.mfx.FrameInfo.CropW = MSDK_ALIGN16(Width);
    pEncode->mfxEncParams.mfx.FrameInfo.CropH = MSDK_ALIGN16(Height);
    // JPEG encoder settings overlap with other encoders settings in mfxInfoMFX structure
    if (MFX_CODEC_JPEG == CodecId)
    {
        pEncode->mfxEncParams.mfx.Interleaved = 1;
        pEncode->mfxEncParams.mfx.Quality = 0;
        pEncode->mfxEncParams.mfx.RestartInterval = 0;
        MSDK_ZERO_MEMORY(pEncode->mfxEncParams.mfx.reserved5);
    }
    pEncode->mfxEncParams.AsyncDepth = 4;
    return MFX_ERR_NONE;
}

mfxStatus CreateAllocator(mfxDHEncode* pEncode)
{ 
    mfxHDL hdl = NULL;
    mfxStatus sts = MFX_ERR_NONE;
    mfxHandleType hdl_t = MFX_HANDLE_D3D11_DEVICE;
    pEncode->memType = D3D9_MEMORY;
    pEncode->hwdev = new CD3D9Device();
    sts = pEncode->hwdev->Init(0,0,MSDKAdapter::GetNumber(pEncode->mfxSession));
    //if(sts<MFX_ERR_NONE) 
        pEncode->memType = SYSTEM_MEMORY;
    hdl_t = MFX_HANDLE_D3D9_DEVICE_MANAGER;
    if (D3D9_MEMORY == pEncode->memType)
    {
        mfxIMPL impl = 0;
        sts = pEncode->hwdev->GetHandle(hdl_t, &hdl);
        MSDK_CHECK_STATUS(sts, "m_hwdev->GetHandle failed");
        pEncode->mfxSession.QueryIMPL(&impl);
        if (impl != MFX_IMPL_SOFTWARE)
        {
            sts = pEncode->mfxSession.SetHandle(hdl_t, hdl);
            MSDK_CHECK_STATUS(sts, "m_mfxSession.SetHandle failed");
        }
        pEncode->pMFXAllocator = new D3DFrameAllocator;
        MSDK_CHECK_POINTER(pEncode->pMFXAllocator, MFX_ERR_MEMORY_ALLOC);

        D3DAllocatorParams *pd3dAllocParams = new D3DAllocatorParams;
        MSDK_CHECK_POINTER(pd3dAllocParams, MFX_ERR_MEMORY_ALLOC);
        pd3dAllocParams->pManager = reinterpret_cast<IDirect3DDeviceManager9 *>(hdl);

        pEncode->pmfxAllocatorParams = pd3dAllocParams;
        sts = pEncode->mfxSession.SetFrameAllocator(pEncode->pMFXAllocator);
        MSDK_CHECK_STATUS(sts, "m_mfxSession.SetFrameAllocator failed");
    }
    else
    {
         pEncode->pMFXAllocator = new SysMemFrameAllocator();
         MSDK_CHECK_POINTER(pEncode->pMFXAllocator, MFX_ERR_MEMORY_ALLOC);
    }
    //initialize memory allocator
    sts = pEncode->pMFXAllocator->Init(pEncode->pmfxAllocatorParams);
    MSDK_CHECK_STATUS(sts, "m_pMFXAllocator->Init failed");

    return MFX_ERR_NONE;
}

mfxStatus ResetMFXComponents(mfxDHEncode* pEncode)
{
    mfxStatus sts = MFX_ERR_NONE;
    sts = pEncode->pmfxENC->Close();
    MSDK_IGNORE_MFX_STS(sts, MFX_ERR_NOT_INITIALIZED);
    MSDK_CHECK_STATUS(sts, "m_pmfxENC->Close failed");

    // free allocated frames
    MSDK_SAFE_DELETE_ARRAY(pEncode->pEncSurfaces);
    if (pEncode->pMFXAllocator)
        pEncode->pMFXAllocator->Free(pEncode->pMFXAllocator->pthis, &pEncode->mEncResponse);

    //allocated frames
    mfxFrameAllocRequest mEncRequest;
    MSDK_ZERO_MEMORY(mEncRequest);
    sts = pEncode->pmfxENC->Query(&pEncode->mfxEncParams, &pEncode->mfxEncParams);
    MSDK_CHECK_STATUS(sts, "Query (for encoder) failed");
    sts = pEncode->pmfxENC->QueryIOSurf(&pEncode->mfxEncParams, &mEncRequest);
    MSDK_CHECK_STATUS(sts, "QueryIOSurf (for encoder) failed");
    if (mEncRequest.NumFrameSuggested < pEncode->mfxEncParams.AsyncDepth)
        return MFX_ERR_MEMORY_ALLOC;
    mEncRequest.NumFrameMin = mEncRequest.NumFrameSuggested;
    MSDK_MEMCPY_VAR(mEncRequest.Info, &(pEncode->mfxEncParams.mfx.FrameInfo), sizeof(mfxFrameInfo));
    sts = pEncode->pMFXAllocator->Alloc(pEncode->pMFXAllocator->pthis, &mEncRequest, &pEncode->mEncResponse);
    MSDK_CHECK_STATUS(sts, "m_pMFXAllocator->Alloc failed");
    pEncode->pEncSurfaces = new mfxFrameSurface1 [pEncode->mEncResponse.NumFrameActual];
    MSDK_CHECK_POINTER(pEncode->pEncSurfaces, MFX_ERR_MEMORY_ALLOC);
    for (int i = 0; i < pEncode->mEncResponse.NumFrameActual; i++)
    {
        memset(&(pEncode->pEncSurfaces[i]), 0, sizeof(mfxFrameSurface1));
        MSDK_MEMCPY_VAR(pEncode->pEncSurfaces[i].Info, &(pEncode->mfxEncParams.mfx.FrameInfo), sizeof(mfxFrameInfo));
        if (pEncode->memType == D3D9_MEMORY || pEncode->memType == D3D11_MEMORY){
            pEncode->pEncSurfaces[i].Data.MemId = pEncode->mEncResponse.mids[i];
        }
        else{
            sts = pEncode->pMFXAllocator->Lock(pEncode->pMFXAllocator->pthis, pEncode->mEncResponse.mids[i], &(pEncode->pEncSurfaces[i].Data));
            MSDK_CHECK_STATUS(sts, "m_pMFXAllocator->Lock failed");
        }
    }
    sts = pEncode->pmfxENC->Init(&pEncode->mfxEncParams);
    return MFX_ERR_NONE;
}

mfxStatus FillBuffers(mfxDHEncode* pEncode, unsigned char* image_data, int image_len,int index)
{
    mfxStatus sts = MFX_ERR_NONE;
    mfxFrameSurface1* surface = &pEncode->pEncSurfaces[index];
    if (D3D11_MEMORY == pEncode->memType || D3D9_MEMORY == pEncode->memType){
        sts = pEncode->pMFXAllocator->Lock(pEncode->pMFXAllocator->pthis, surface->Data.MemId, &surface->Data);
        MSDK_CHECK_STATUS(sts, "m_pMFXAllocator->Lock failed");
    }

    if(AV_PIX_FMT_GRAY8==g_format){
        memcpy(pEncode->swframe->data[0],image_data,image_len);
        memset(pEncode->swframe->data[1],128,image_len/2);
    }
    else{
        if(av_image_fill_arrays(pEncode->frame->data,pEncode->frame->linesize, image_data, g_format, 
            pEncode->frame->width, pEncode->frame->height,16)<0)
            av_image_fill_arrays(pEncode->frame->data,pEncode->frame->linesize, image_data, g_format, 
                pEncode->frame->width, pEncode->frame->height,1);
        switch(g_format){
            case AV_PIX_FMT_ARGB:
                libyuv::ARGBToI420(pEncode->frame->data[0],pEncode->frame->linesize[0],
                    pEncode->swframe->data[0],pEncode->swframe->linesize[0],
                    pEncode->swframe->data[2],pEncode->swframe->linesize[2],
                    pEncode->swframe->data[1],pEncode->swframe->linesize[1],
                    pEncode->frame->width,pEncode->frame->height);
                break;
            default:
                sws_scale(pEncode->sws,(const uint8_t **)pEncode->frame->data,pEncode->frame->linesize,
                    0,pEncode->frame->height, pEncode->swframe->data,pEncode->swframe->linesize);
                break;
        }
    }
    memcpy(surface->Data.Y,pEncode->swframe->data[0],
        pEncode->swframe->width*pEncode->swframe->height);
    memcpy(surface->Data.U,pEncode->swframe->data[1],
        pEncode->swframe->width*pEncode->swframe->height/2);
    if (D3D11_MEMORY == pEncode->memType || D3D9_MEMORY == pEncode->memType){
        sts = pEncode->pMFXAllocator->Unlock(pEncode->pMFXAllocator->pthis, surface->Data.MemId, &surface->Data);
        MSDK_CHECK_STATUS(sts, "m_pMFXAllocator->Unlock failed");
    }
    MSDK_CHECK_STATUS(sts, "FillBuffers failed");
    return sts;
}

mfxStatus SyncEncode(mfxDHEncode* pEncode)
{
    if(pEncode->mfxEncSyncP){
        mfxStatus sts = pEncode->mfxSession.SyncOperation(pEncode->mfxEncSyncP,MSDK_ENC_WAIT_INTERVAL);
        if (MFX_ERR_NONE == sts){
            pEncode->encode(pEncode->user_data,pEncode->mfxBS.Data,(int&)pEncode->mfxBS.DataLength);
            pEncode->mfxBS.DataLength = 0;
            pEncode->mfxBS.DataOffset = 0;
            pEncode->mfxEncSyncP = nullptr;
        }
        else if(sts == MFX_ERR_GPU_HANG){
            sts = pEncode->mfxSession.SyncOperation(pEncode->mfxEncSyncP, 0);
            MSDK_CHECK_STATUS(sts, "SyncOperation failed");
        }
    }
    return MFX_ERR_NONE;
}
void MfxBeginEncode(void** handle, void* obj, int width, int height, int framerate, __int64 bitrate, int code, void(*encode)(void* obj, unsigned char* data, int& len))
{ 
    mfxStatus sts = MFX_ERR_NONE;
    *handle = (void*)new mfxDHEncode();
    mfxDHEncode* pEncode = (mfxDHEncode*)*handle;
    pEncode->user_data = obj;
    pEncode->encode = encode;
    InitSession(pEncode->mfxSession);
    if(!pEncode->mfxSession) return;
    pEncode->pmfxENC = new MFXVideoENCODE(pEncode->mfxSession);
    if(!pEncode->pmfxENC) return;

    sts = CreateAllocator(pEncode);
    if(sts < MFX_ERR_NONE) return;
    sts = InitMfxEncParams(pEncode,code==27?MFX_CODEC_AVC:MFX_CODEC_HEVC,width,height,framerate,bitrate);
    if(sts < MFX_ERR_NONE) return;
    sts = ResetMFXComponents(pEncode);
    if(sts < MFX_ERR_NONE) return;

    pEncode->frame = av_frame_alloc();
    pEncode->swframe = av_frame_alloc();
    pEncode->frame->format = g_format;
    pEncode->frame->width  = width;
    pEncode->frame->height = height; 
    pEncode->swframe->format = AV_PIX_FMT_YUV420P;
    pEncode->swframe->width  = width;
    pEncode->swframe->height = height;
    pEncode->frame->pts = pEncode->swframe->pts = 0;
    if(av_frame_get_buffer(pEncode->frame, 16)<0)
        av_frame_get_buffer(pEncode->frame, 1);
    if(av_frame_get_buffer(pEncode->swframe, 16))
        av_frame_get_buffer(pEncode->swframe, 1);
    pEncode->sws = sws_getContext(width, height, g_format,width, height, g_format,SWS_FAST_BILINEAR, NULL, NULL, NULL);
}

void MfxImageEncode(void** handle, unsigned char* image_data, int image_lens)
{
    mfxStatus sts = MFX_ERR_NONE;
    mfxDHEncode* pEncode = (mfxDHEncode*)*handle;
    
    mfxU16 nEncSurfIdx = 0;
    while(pEncode->pEncSurfaces[nEncSurfIdx].Data.Locked){
        nEncSurfIdx++;
        if(nEncSurfIdx==pEncode->mEncResponse.NumFrameActual){
            nEncSurfIdx = 0;
            boost::this_thread::sleep(boost::posix_time::millisec(10));
        }
    }

#ifdef TEST
    test_start();
#endif
    FillBuffers(pEncode,image_data,image_lens,nEncSurfIdx);
#ifdef TEST
    double tt = test_check();
#endif
    for (;;){
        pEncode->encCtrl.FrameType = MFX_FRAMETYPE_I | MFX_FRAMETYPE_IDR | MFX_FRAMETYPE_REF;
        sts = pEncode->pmfxENC->EncodeFrameAsync(&pEncode->encCtrl, 
            &pEncode->pEncSurfaces[nEncSurfIdx], &pEncode->mfxBS, &pEncode->mfxEncSyncP);
        SyncEncode(pEncode);
        if (MFX_ERR_NONE < sts && !pEncode->mfxEncSyncP) // repeat the call if warning and no output
        {
            if (MFX_WRN_DEVICE_BUSY == sts){
                MSDK_SLEEP(1); // wait if device is busy
            }
            else{
                sts = MFX_ERR_NONE;
                break;
            }
        }
        else if(MFX_ERR_NOT_ENOUGH_BUFFER == sts)
        {
            mfxVideoParam par;
            MSDK_ZERO_MEMORY(par);
            sts = pEncode->pmfxENC->GetVideoParam(&par);
            if(sts<MFX_ERR_NONE) break;
            sts = ExtendMfxBitstream(&pEncode->mfxBS, par.mfx.BufferSizeInKB * 1000);
            if(sts<MFX_ERR_NONE) WipeMfxBitstream(&pEncode->mfxBS);
        }
        else
        {
            // get new task for 2nd bitstream in ViewOutput mode
            MSDK_IGNORE_MFX_STS(sts, MFX_ERR_MORE_BITSTREAM);
            break;
        }
    }
    if((MFX_ERR_DEVICE_LOST == sts || MFX_ERR_DEVICE_FAILED == sts) && pEncode->hwdev) {
        sts = pEncode->hwdev->Reset();
        if(sts<MFX_ERR_NONE) return;
        ResetMFXComponents(pEncode);
        if(sts<MFX_ERR_NONE) return;
    }

#ifdef TEST
    printf("%f-%f ms\r\n",tt,test_check());
#endif
}
void MfxEndEncode(void** handle)
{
    mfxDHEncode* pEncode = (mfxDHEncode*)*handle;
    mfxStatus sts = MFX_ERR_NONE;
    while (MFX_ERR_NONE <= sts)
    {
        for (;;)
        {
            pEncode->encCtrl.FrameType = MFX_FRAMETYPE_I | MFX_FRAMETYPE_IDR | MFX_FRAMETYPE_REF;
            sts = pEncode->pmfxENC->EncodeFrameAsync(&pEncode->encCtrl, nullptr, &pEncode->mfxBS, &pEncode->mfxEncSyncP);
            SyncEncode(pEncode);
            if (MFX_ERR_NONE < sts && !pEncode->mfxEncSyncP) // repeat the call if warning and no output
            {
                if (MFX_WRN_DEVICE_BUSY == sts){
                    MSDK_SLEEP(1); // wait if device is busy
                }
                else{
                    sts = MFX_ERR_NONE;
                    break;
                }
            }
            else
            {
                // get new task for 2nd bitstream in ViewOutput mode
                MSDK_IGNORE_MFX_STS(sts, MFX_ERR_MORE_BITSTREAM);
                break;
            }
        }
    }
    if(pEncode){
        MSDK_SAFE_DELETE(pEncode->pmfxENC);
        MSDK_SAFE_DELETE_ARRAY(pEncode->pEncSurfaces);
        if (pEncode->pMFXAllocator){
            pEncode->pMFXAllocator->Free(pEncode->pMFXAllocator->pthis, &pEncode->mEncResponse);
        }
        MSDK_SAFE_DELETE(pEncode->pMFXAllocator);
        MSDK_SAFE_DELETE(pEncode->pmfxAllocatorParams);
        MSDK_SAFE_DELETE(pEncode);
    }
}